#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time : 2021/1/22 12:13
# @Author : lijunpeng
# @File keyword_search

import ahocorasick
import os
import datetime, time

# 排除的后缀
def Type_of_file(file):
    file_exts = ['.png', '.jpg', '.jpeg', '.gif', '.mp4', '.mp3', '.svg', '.eot', '.font', '.ico']
    file_ext = os.path.splitext(file)[-1]
    if file_ext in file_exts:
        return True
    else:
        return False

# 遍历文件
def Traverse_file(path):
    files = os.listdir(path)
    for file in files:
        file_d = os.path.join(path, file)
        if os.path.isdir(file_d):
            Traverse_file(file_d)
        elif Type_of_file(file_d):
            break
        else:
            Search(file_d)

# 自动机匹配文件
def Search(file):
    A = ahocorasick.Automaton()
    # 向trie树中添加关键词
    with open("keyword.txt", 'r', encoding='utf-8') as fp:
        while True:
            tok = fp.readline().strip("\n")
            if len(tok) < 1:
                break
            else:
                A.add_word(str(tok), (1, str(tok)))
    # 将trie树转化为自动机
    A.make_automaton()
    g1 = open("result/result.txt", 'a+', encoding='utf-8')
    for count, line in enumerate(open(file, 'rb')):
        for k, (i,t) in A.iter(line.decode('utf-8').strip()):
            print("文件:{}在第{}行存在违禁词:{}".format(file, count+1, t))
            g1.write("文件:{}在第{}行存在违禁词:{}\n".format(file, count+1, t))
    g1.close()


if __name__ == "__main__":
    filepath = input("请输入要搜索的路径：")
    if os.path.isdir(filepath):
        g1 = open("result/result.txt", 'a+', encoding='utf-8')
        g1.write("查询时间：" + str(datetime.datetime.now()) + "\n")
        g1.close()
        t1 = time.time()
        try:
            Traverse_file(filepath)
            t2 = time.time()
            print("扫描完毕，用时", t2-t1)
        except Exception as e:
            print(e)
    else:
        print("请输入正确得路径！")
